import numpy as np
import matplotlib.pyplot as plt
import os
import pandas as pd
import nibabel as nib
from skimage.measure import marching_cubes

def visualize_3d_patch(ct_path, mask_path, metadata, save_dir=None, public_id=None, label_id=None):
    """
    正确可视化的3D CT块和骨折掩码，解决方位定位问题
    
    参数:
    ct_path -- CT影像块的.npy文件路径
    mask_path -- 骨折掩码块的.npy文件路径
    metadata -- 包含坐标信息的元数据DataFrame
    save_dir -- 保存图像的可选目录
    public_id -- 样本ID用于标题
    label_id -- 骨折ID用于标题
    """
    # 获取元数据中的原始坐标
    sample_meta = metadata[(metadata['public_id'] == public_id) & 
                          (metadata['label_id'] == label_id)]
    if sample_meta.empty:
        print(f"找不到元数据: public_id={public_id}, label_id={label_id}")
        return
    
    orig_x = sample_meta['center_x'].values[0]
    orig_y = sample_meta['center_y'].values[0]
    orig_z = sample_meta['center_z'].values[0]
    
    ct_block = np.load(ct_path)
    mask_block = np.load(mask_path)
    
    if ct_block.shape != mask_block.shape:
        print(f"CT和掩码形状不匹配: CT {ct_block.shape} vs Mask {mask_block.shape}")
        return
    height, width,depth = ct_block.shape
    
    height, width,depth = ct_block.shape
    
    height, width,depth = ct_block.shape
    
    fig = plt.figure(figsize=(18, 12))
    fig.suptitle(f'Fracture Visualization | ID: {public_id}-{label_id}\n'
                f'Original Coords: X={orig_x}, Y={orig_y}, Z={orig_z}', 
                fontsize=16)
    ax1 = fig.add_subplot(231)
    z_slice = depth // 2 


    axial_img = np.rot90(ct_block[:, :, z_slice], k=1)
    
    im1 = ax1.imshow(axial_img, cmap='gray', origin='lower',
                     extent=[0, height, 0, width])
    
    mask_slice = np.rot90(mask_block[:, :, z_slice], k=1)
    ax1.contour(mask_slice, levels=[0.5], colors='red', linewidths=1)
    
    ax1.set_title(f'Axial Slice (Z={z_slice}/{depth})')
    ax1.set_xlabel('Y (Anterior-Posterior)')
    ax1.set_ylabel('X (Left-Right)')
    ax1.grid(False)
    
    ax2 = fig.add_subplot(232)
    y_slice = height // 2
    ax2 = fig.add_subplot(232)
    y_slice = height // 2
    ax2 = fig.add_subplot(232)
    y_slice = height // 2
    coronal_img = ct_block[:, y_slice, :].T
    
    im2 = ax2.imshow(coronal_img, cmap='gray', origin='lower',
                     extent=[0, depth, 0, width])
    
    mask_slice = mask_block[:, y_slice, :].T
    ax2.contour(mask_slice, levels=[0.5], colors='red', linewidths=1)
    
    ax2.set_title(f'Coronal Slice (Y={y_slice}/{height})')
    ax2.set_xlabel('Z (Inferior-Superior)')
    ax2.set_ylabel('X (Left-Right)')
    ax2.grid(False)
    ax3 = fig.add_subplot(233)
    x_slice = width // 2
    ax3 = fig.add_subplot(233)
    x_slice = width // 2

    ax3 = fig.add_subplot(233)
    x_slice = width // 2

    sagittal_img = ct_block[x_slice, :,: ].T
    
    im3 = ax3.imshow(sagittal_img, cmap='gray', origin='lower',
                     extent=[0, depth, 0, height])
    mask_slice = mask_block[x_slice, :, :].T
    ax3.contour(mask_slice, levels=[0.5], colors='red', linewidths=1)
    
    ax3.set_title(f'Sagittal Slice (X={x_slice}/{width})')
    ax3.set_xlabel('Z (Inferior-Superior)')
    ax3.set_ylabel('Y (Anterior-Posterior)')
    ax3.grid(False)
    

    ax4 = fig.add_subplot(234, projection='3d')

    if np.any(mask_block > 0):
        try:
            verts, faces, _, _ = marching_cubes(mask_block, level=0.5, step_size=1)
            ax4.plot_trisurf(verts[:, 1], verts[:, 0], faces, verts[:, 2],
                            cmap='Spectral', lw=0.5, antialiased=True, alpha=0.7)
        
            ax4.scatter([height//2], [width//2], [depth//2], c='cyan', s=100, marker='*')
            ax4.text(height//2, width//2, depth//2, ' Center', color='cyan', fontsize=10)
            
        except Exception as e:
            print(f"3D渲染失败: {e}")
    else:
        ax4.text(0.3, 0.5, 0.5, "No Fracture Detected", fontsize=12)
    
    ax4.set_title('3D Fracture Surface')
    ax4.set_xlabel('Y (A-P)')
    ax4.set_ylabel('X (L-R)')
    ax4.set_zlabel('Z (I-S)')
    ax4.view_init(elev=30, azim=45)


    ax5 = fig.add_subplot(235)
    ax5.hist(ct_block.ravel(), bins=100, color='blue', alpha=0.7)
    
    ax5 = fig.add_subplot(235)
    ax5.hist(ct_block.ravel(), bins=100, color='blue', alpha=0.7)

    ax5 = fig.add_subplot(235)
    ax5.hist(ct_block.ravel(), bins=100, color='blue', alpha=0.7)
    
    ax5.axvline(x=-1000, color='gray', linestyle='--', alpha=0.5)
    ax5.axvline(x=-100, color='gray', linestyle='--', alpha=0.5)
    ax5.axvline(x=0, color='gray', linestyle='--', alpha=0.5)
    ax5.axvline(x=40, color='gray', linestyle='--', alpha=0.5)
    ax5.axvline(x=400, color='gray', linestyle='--', alpha=0.5)
    ax5.axvline(x=1000, color='gray', linestyle='--', alpha=0.5)
    
    ax5.text(-950, ax5.get_ylim()[1]*0.9, 'Air', ha='center')
    ax5.text(-500, ax5.get_ylim()[1]*0.85, 'Lung', ha='center')
    ax5.text(-50, ax5.get_ylim()[1]*0.9, 'Fat', ha='center')
    ax5.text(20, ax5.get_ylim()[1]*0.85, 'Water', ha='center')
    ax5.text(200, ax5.get_ylim()[1]*0.9, 'Soft Tissue', ha='center')
    ax5.text(700, ax5.get_ylim()[1]*0.85, 'Bone', ha='center')
    
    ax5.set_title('CT Value Distribution (HU)')
    ax5.set_xlabel('Hounsfield Units')
    ax5.set_ylabel('Frequency')
    ax5.set_xlim([-1024, 2000])
    ax5.grid(True, alpha=0.3)
    ax6 = fig.add_subplot(236)
    ax6.axis('off')
    
    fracture_types = {
        0: "Displaced",
        1: "Non-displaced",
        2: "Buckle",
        3: "Segmental"
    }
    class_label = sample_meta['class_label'].values[0]
    frac_type = fracture_types.get(class_label, "Unknown")
    
    voxel_count = np.sum(mask_block)
    hu_min = np.min(ct_block)
    hu_max = np.max(ct_block)
    hu_mean = np.mean(ct_block)
    
    info_text = (
        f"⚫ Fracture Type: {frac_type} (Class {class_label})\n"
        f"⚫ Voxel Count: {voxel_count} voxels\n"
        f"⚫ Volume: {voxel_count:.1f} mm³\n"
        f"⚫ HU Range: [{hu_min:.1f}, {hu_max:.1f}]\n"
        f"⚫ Mean HU: {hu_mean:.1f}\n"
        f"⚫ Patch Dimensions: {depth}×{height}×{width}"
    )
    
    ax6.text(0.05, 0.5, info_text, fontsize=12, 
             bbox=dict(facecolor='lightblue', alpha=0.3, pad=15))
    
    plt.tight_layout(rect=[0, 0, 1, 0.95])
    
    if save_dir:
        os.makedirs(save_dir, exist_ok=True)
        save_path = os.path.join(save_dir, f'fracture_{public_id}_{label_id}.png')
        plt.savefig(save_path, dpi=150, bbox_inches='tight')
        print(f"已保存可视化图像到: {save_path}")
        plt.close()
    else:
        plt.show()

if __name__ == "__main__":
    # # 查看刚划分好的数据块
    # output_dir = './ribfrac-dataset/train'
    # metadata = pd.read_csv(os.path.join(output_dir, 'metadata', 'metadata.csv'))
    # sample = metadata.iloc[0]
    # ct_path = os.path.join(output_dir, 'images', f"{sample['public_id']}_{sample['label_id']}_ct.npy")
    # mask_path = os.path.join(output_dir, 'labels', f"{sample['public_id']}_{sample['label_id']}_mask.npy")
    # visualize_3d_patch(
    #     ct_path=ct_path,
    #     mask_path=mask_path,
    #     metadata=metadata,
    #     save_dir='./ribfrac-patches/visualizations',
    #     public_id=sample['public_id'],
    #     label_id=sample['label_id']
    # )


    data_test_dir = './ribfrac-dataset/test'
    # 更换可视化资源修改此处
    predict_path = './segment_test_results/3dunet' 
    
    metadata = pd.read_csv(os.path.join(data_test_dir, 'metadata', 'metadata.csv'))
    sample = metadata.iloc[1]
    ct_path_gt = os.path.join(data_test_dir, 'images', f"{sample['public_id']}_{sample['label_id']}_ct.npy")
    mask_path_gt = os.path.join(data_test_dir, 'labels', f"{sample['public_id']}_{sample['label_id']}_mask.npy")
    
    visualize_3d_patch(
        ct_path=ct_path_gt,
        mask_path=mask_path_gt,
        metadata=metadata,
        save_dir=os.path.join(predict_path, 'visualizations_gt'),
        public_id=sample['public_id'],
        label_id=sample['label_id']
    )
    ct_path_predict = os.path.join(data_test_dir, 'images', f"{sample['public_id']}_{sample['label_id']}_ct.npy")
    mask_path_predict = os.path.join(predict_path, 'pred', f"{sample['public_id']}_{sample['label_id']}_binary.npy")
    visualize_3d_patch(
        ct_path=ct_path_predict,
        mask_path=mask_path_predict,
        metadata=metadata,
        save_dir=os.path.join(predict_path, 'visualizations_predict'),
        public_id=sample['public_id'],
        label_id=sample['label_id']
    )

